#!/usr/bin/env python
# coding: utf-8

# 数据载入
from keras.layers import Embedding, Reshape, Input, Dot
import keras.backend as K
from keras import Model
import pandas as pd
import numpy as np
rating = pd.read_csv(r"./ratings.dat", header=None,
                     names=['UserIDs', 'MovieIDs', 'Ratings', 'Timestamp'], sep="::")
num_user = np.max(rating["UserIDs"])
num_movie = np.max(rating["MovieIDs"])
print(num_user, num_movie, len(rating))
# 模型搭建
K.clear_session()

def Recmand_model(num_user, num_movie, k):
    input_uer = Input(shape=[None, ], dtype="int32")
    model_uer = Embedding(num_user+1, k, input_length=1)(input_uer)
    model_uer = Reshape((k,))(model_uer)
    input_movie = Input(shape=[None, ], dtype="int32")
    model_movie = Embedding(num_movie+1, k, input_length=1)(input_movie)
    model_movie = Reshape((k,))(model_movie)
    out = Dot(1)([model_uer, model_movie])
    model = Model(inputs=[input_uer, input_movie], outputs=out)
    model.compile(loss='mse', optimizer='Adam')
    model.summary()
    return model

model = Recmand_model(num_user, num_movie, 100)

# 数据准备
user_train = rating["UserIDs"].values
movie_train = rating["MovieIDs"].values
rating_train = rating["Ratings"].values

user_pre = user_train[10000:11250]  # 测试集
movie_pre = movie_train[10000:11250]  # 测试集
rating_pre = rating_train[10000:11250]

user_train = user_train[0:10000]  # 训练集
movie_train = movie_train[0:10000]  # 训练集

train_x = [user_train, movie_train]
train_y = rating["Ratings"].values

train_y = train_y[0:10000]  # 训练集

# 模型训练
model.fit(train_x, train_y, batch_size=100, epochs=10)

# 预测评估
MSE = 0
for i in range(int(len(user_pre))):
    predict = model.predict(
        [np.array([user_pre[1]]), np.array([movie_pre[1]])])
    rating_true = rating_pre[i]
    MSE += (predict[0][0] - rating_true)**2
MSE = MSE/len(user_pre)
print(MSE)
